package com.lap.auth.filter;

import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletRequestWrapper;
import java.util.*;

public class HeaderModifyRequestWrapper extends HttpServletRequestWrapper {

  private final Map<String, String> customHeaders;
  private final Map<String, String[]> customParameters;

  public HeaderModifyRequestWrapper(HttpServletRequest request) {
    super(request);
    this.customHeaders = new HashMap<>();
    this.customParameters = new HashMap<>();
  }

  // 头信息相关方法
  public void putHeader(String name, String value) {
    this.customHeaders.put(name.toLowerCase(), value);
  }

  public void removeHeader(String name) {
    this.customHeaders.remove(name.toLowerCase());
  }

  @Override
  public String getHeader(String name) {
    String headerValue = customHeaders.get(name.toLowerCase());
    if (headerValue != null) {
      return headerValue;
    }
    return super.getHeader(name);
  }

  @Override
  public Enumeration<String> getHeaderNames() {
    Set<String> set = new HashSet<>(customHeaders.keySet());
    Enumeration<String> e = super.getHeaderNames();
    while (e.hasMoreElements()) {
      set.add(e.nextElement());
    }
    return Collections.enumeration(set);
  }

  @Override
  public Enumeration<String> getHeaders(String name) {
    List<String> values = new ArrayList<>();

    String customValue = customHeaders.get(name.toLowerCase());
    if (customValue != null) {
      values.add(customValue);
    }

    Enumeration<String> e = super.getHeaders(name);
    while (e.hasMoreElements()) {
      values.add(e.nextElement());
    }

    return Collections.enumeration(values);
  }

  // 参数相关方法（可选）
  public void putParameter(String name, String value) {
    this.customParameters.put(name, new String[] {value});
  }

  @Override
  public String getParameter(String name) {
    if (customParameters.containsKey(name)) {
      String[] values = customParameters.get(name);
      return values != null && values.length > 0 ? values[0] : null;
    }
    return super.getParameter(name);
  }

  @Override
  public Map<String, String[]> getParameterMap() {
    Map<String, String[]> map = new HashMap<>(super.getParameterMap());
    map.putAll(customParameters);
    return Collections.unmodifiableMap(map);
  }

  @Override
  public Enumeration<String> getParameterNames() {
    Set<String> set = new HashSet<>(customParameters.keySet());
    Enumeration<String> e = super.getParameterNames();
    while (e.hasMoreElements()) {
      set.add(e.nextElement());
    }
    return Collections.enumeration(set);
  }

  @Override
  public String[] getParameterValues(String name) {
    if (customParameters.containsKey(name)) {
      return customParameters.get(name);
    }
    return super.getParameterValues(name);
  }
}
